
from transformers import Wav2Vec2Processor, Wav2Vec2FeatureExtractor
from wav2vec import Wav2Vec2Model, Wav2Vec2ForSpeechClassification
from transformers import AutoTokenizer
from onnxruntime import InferenceSession
import torch
import os
def main():
    modelPath = './checkpoint/et/'
    outPath = os.path.join(modelPath,"onnx")
    
    wav2vec2_path = "jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn"
    pt_model  = Wav2Vec2Model.from_pretrained(modelPath + wav2vec2_path)        
    # tokenizer = AutoTokenizer.from_pretrained(modelPath + wav2vec2_path)
    # tokenizer.save_pretrained(outPath)
    # print('tokenizer saved')
    # pt_model.save_pretrained(outPath)
    # print('pt_model saved')
    audio_len = 250000

    onnx_model_name = os.path.join(modelPath ,'onnx/model.onnx')
    x = torch.randn(1, audio_len, requires_grad=True)
    torch.onnx.export(pt_model,                        # model being run
                        x,                              # model input (or a tuple for multiple inputs)
                        onnx_model_name,                # where to save the model (can be a file or file-like object)
                        export_params=True,             # store the trained parameter weights inside the model file
                        opset_version=11,               # the ONNX version to export the model to
                        do_constant_folding=True,       # whether to execute constant folding for optimization
                        input_names = ['input'],        # the model's input names
                        output_names = ['output'],      # the model's output names
                        dynamic_axes={'input' : {1 : 'audio_len'},    # variable length axes
                                    'output' : {1 : 'audio_len'}})

    # session = InferenceSession("onnx/model.onnx")
    # # ONNX Runtime expects NumPy arrays as input
    # inputs = tokenizer("Using DistilBERT with ONNX Runtime!", return_tensors="np")
    # outputs = session.run(output_names=["last_hidden_state"], input_feed=dict(inputs))






if __name__ == '__main__':
    main()